package com.example.deepseek_ai.controller;

import jakarta.annotation.PostConstruct;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.File;
import java.util.ArrayList;
import java.util.List;

@Slf4j
@RestController
@RequestMapping("/rag")
@AllArgsConstructor
public class RagDemoController {
    private final VectorStore vectorStore;

    @PostConstruct
    public void init() {
        addVectorStore();  // 启动时执行加载
    }
    public void addVectorStore() {
        try {
            // 你可以将多个文件放在 /resources/documents/ 目录下
            List<String> filenames = List.of(
                    "data2025.xlsx",
                    "清理专业表.xlsx",
                    "data2017-2024.xlsx"
            );

            TokenTextSplitter splitter = new TokenTextSplitter(500, 200, 1, 8192, true);
            List<Document> allDocs = new ArrayList<>();

            for (String filename : filenames) {
                Resource resource = new ClassPathResource("documents/" + filename);
                TikaDocumentReader reader = new TikaDocumentReader(resource);
                List<Document> docs = splitter.apply(reader.get());
                allDocs.addAll(docs);
                log.info("已处理文件 {}，生成文档片段数：{}", filename, docs.size());
            }

            vectorStore.add(allDocs);
            log.info("共成功加载 {} 条文档片段", allDocs.size());

        } catch (Exception e) {
            log.error("加载文档失败", e);
        }
    }

}
